from polycoherence import polycoherence,plot_polycoherence
import matplotlib.pyplot as plt
from scipy import signal
import numpy as np
import samplerate
import librosa
import keras

def plot_signal(audio_data, title=None):
    plt.figure(figsize=(12, 3.5), dpi=300)
    plt.plot(audio_data, linewidth=1)
    plt.title(title,fontsize = 16)
    plt.tick_params(labelsize=12)
    plt.grid()
    plt.show()

def band_pass_filter(original_signal, order, fc1,fc2, fs):
    b, a = signal.butter(N=order, Wn=[2*fc1/fs,2*fc2/fs], btype='bandpass')
    new_signal = signal.lfilter(b, a, original_signal)
    return new_signal

audio_path = './a0001.wav'

audio_data, fs = librosa.load(audio_path, sr=2000)
plot_signal(audio_data, title='Init Filter')
print(type(audio_data))
print("原始音频数据点数：", audio_data.shape, "采样率：", fs)

audio_data = band_pass_filter(audio_data, 2, 25, 400, fs)
plot_signal(audio_data, title='After Filter')
print(len(audio_data))

down_sample_audio_data = samplerate.resample(audio_data.T, 1000 / fs, converter_type='sinc_best').T
plot_signal(down_sample_audio_data, title='Down_sampled')
print(len(down_sample_audio_data))

down_sample_audio_data = down_sample_audio_data / np.max(np.abs(down_sample_audio_data))
plot_signal(down_sample_audio_data, title='Normalized')
print(len(down_sample_audio_data))

ex_audio_data = down_sample_audio_data
freq1, freq2, bi_spectrum = polycoherence(
       ex_audio_data,
       nfft=1024,
       nperseg=256,
       noverlap = 100,
       fs = 1000,
       norm=None)
print(ex_audio_data.shape)
bi_spectrum = np.array(abs(bi_spectrum))  # calculate bi_spectrum
bi_spectrum = 255 * (bi_spectrum - np.min(bi_spectrum)) / (np.max(bi_spectrum) - np.min(bi_spectrum))
print(bi_spectrum,bi_spectrum.shape)
plot_polycoherence(freq1, freq2, bi_spectrum)
# 修改尺寸以便于投入神经网络
bi_spectrum = bi_spectrum.reshape((1, 256, 256))
print(bi_spectrum.shape)

model = keras.models.load_model("./model_file/h5_file/model_4.h5")
print(model.summary())
predictions = model.predict(bi_spectrum.reshape((1, 256, 256, 1)))
print(predictions)
result = np.argmax(predictions)
print(result)

# 'AS', 'MS', 'MR', 'MVP'
print("心音可能异常，希望您早日就医检查")
if result == 0:
    print("诊断结果为：主动脉瓣膜狭窄（AS）")
elif result == 1:
    print("诊断结果为：二尖瓣狭窄（MS）")
elif result == 2:
    print("诊断结果为：二尖瓣反流（MR）")
elif result == 3:
    print("诊断结果为：二尖瓣脱垂（MVP）")